import xlrd

from gensim.models import word2vec
from word_count import split_words_jieba
from base_handle import BaseHandle # 引入工具类

baseHandle = BaseHandle() #实例化

'''
1.2 word2vec 拓展关键词词库
'''

def wordcut_listinlist(url):
    '''读取excel评论并分别对每条评论进行分词操作'''
    excel = xlrd.open_workbook(url)  # 打开excel文件
    table = excel.sheet_by_index(0)  # 根据下标获取工作薄，这里获取第一个
    comments_list = table.col_values(1, start_rowx=1)  # 获取第一列的内容
    list = []
    for i in comments_list:
        a = split_words_jieba(i)
        list.append(a)
    return list  # [['xx','xx]、['X']、['XXX']]


def word2vec_trainmodel(url):  # 训练一次就行了
    '''Word2vec关键词扩展——训练模型'''
    sentennce = wordcut_listinlist(url)
    # 1.建立模型
    model = word2vec.Word2Vec(sentennce, vector_size=200)  # 训练skip-gram模型，默认window=5
    # 2.保存模型，以便重用
    model.save(u"word2vec_last.model")


def similar_words(word):
    '''计算某个词的相关词列表——前提是已存在训练好的模型'''
    model1 = word2vec.Word2Vec.load("word2vec_last.model")  # 模型的加载方式
    y2 = model1.wv.most_similar(word, topn=100)  # 10个最相关的
    print("和{}最相关的词有：".format(word))
    for item in y2:
        print(item[0], '%.3f' % item[1])


if __name__ == "__main__":
    # word2vec_trainmodel(baseHandle.get_file_abspath('语料库_京东_5000条评论.xlsx')) # 5000条评论，运行2分钟
    similar_words('物流')